from transformers import BertTokenizer, BertModel
import pandas as pd
import json
import pickle


class get_vd:
    def __init__(self):
        self.data

    def file_processing(self):
        # 语料处理模型实现对文件的划分
        pass

    def text_embedding(self, text, bert_path = '../bert-base-chinese'):
        tokenizer = BertTokenizer.from_pretrained(bert_path)
        bert = BertModel.from_pretrained(bert_path)
        text_tokenizer = tokenizer(text, padding = True, truncation = True, max_length = 512,
                                   return_tensors = 'pt')
        output = bert(input_ids = text_tokenizer['input_ids'], attention_mask = text_tokenizer['attention_mask'],
                      token_type_ids = text_tokenizer['token_type_ids'],
                      return_dict = False)
        # torch.Size([1, 498, 768])->pooled_output[0]
        # torch.Size([1, 768])->pooled_output[-1]
        return output[-1][0].detach().numpy()  # 均值编码,转化数组

    def create_plk(self, txt_path, save_path = 'data.pkl'):
        '''
        生成novel.txt文件
        前提：需要获取大量小说语料
        '''
        data = pd.read_csv(txt_path)
        text_list = []
        for index, text in enumerate(data):
            dict_ = {}
            dict_['id'] = index
            dict_['text'] = text
            dict_['text_embedding'] = self.text_embedding(text)
            text_list.append(dict_)
        # 保存列表数据到文件
        with open(save_path, 'wb') as f:
            pickle.dump(text_list, f)

    def add_faiss_vd(self,embedding,index_tpye=''):
        '''faiss向量数据库'''
        num_vectors=embedding.shape[0]
        dim=embedding.shape[-1]
        import faiss
        #索引类型判断
        if index_tpye:
            #L2距离索引
            index=faiss.IndexFlatL2(dim)
        else:
            #点积索引
            index=faiss.IndexFlatIP
        index.add(embedding)

    def search_faiss_vd(self,q,k):
        index=faiss.IndexFlatL2(q.shape[0])
        _,I=index.search(q,k)
        return I

if __name__ == '__main__':
    txt = '../data/novel.txt'
    get_vd.create_plk(txt)
